"""
MultiQueryRetriever 用LLM 将用户查询改写成多个查询
ContextualCompressionRRetriever 用LLM对召回文本进行压缩过过滤
EnsembleRetriever 集成多个招呼器
LongContextReorder 将召回内容进行重新排序
MultiVectorRetriever 为每个候选文档保存多个向量表示
ParentDocumentRetriever 小块文本向量当做索引，返回大块文本
SelfQueryRetriever 将自然语言查询转换为结构化查询
TimeWeightedVectorStoreRetriever 召回文本时考虑时效性因素
"""


import os
import dotenv
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_text_splitters import CharacterTextSplitter
from langchain_redis import RedisVectorStore


def fun_1():
    dotenv.load_dotenv()

    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL"),
        base_url=os.getenv("CLOSE_BASE_URL"),
        api_key=os.getenv("CLOSE_API_KEY")
    )

    # 加载TXT文档
    loader = TextLoader(file_path="F:/data/1.txt", encoding="utf-8")
    documents = loader.load()

    # 创建文本分块器
    test_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=10)
    chunks = test_splitter.split_documents(documents)

    #获取向量数据库
    db = FAISS.from_documents(documents=chunks, embedding=embedding_model)

    #创建向量搜索器
    retriever = db.as_retriever()

    response = retriever.invoke("协和")
    print(response)

#fun_1()

#默认使用相似性搜索
def fun_2():
    documents = [
        Document(
            page_content="近日，中国科学家在量子计算领域取得重大突破，实现了100量子比特的纠缠态操控。",
            metadata={"source": "science_news.txt", "category": "科技新闻", "date": "2024-01-15"}
        ),
        Document(
            page_content="《三体》是刘慈欣创作的科幻小说三部曲，包括《三体》、《黑暗森林》和《死神永生》。",
            metadata={"source": "literature_review.txt", "category": "文学作品", "author": "刘慈欣"}
        ),
        Document(
            page_content="人工智能在医疗诊断中的应用越来越广泛，特别是在影像识别方面表现突出。",
            metadata={"source": "ai_research.txt", "category": "学术研究", "field": "医疗AI"}
        ),
        Document(
            page_content="巴黎奥运会筹备工作进入冲刺阶段，预计将有超过200个国家和地区参与。",
            metadata={"source": "sports_news.txt", "category": "体育新闻", "event": "2024巴黎奥运会"}
        ),
        Document(
            page_content="区块链技术不仅应用于加密货币，还在供应链管理、数字身份验证等领域展现潜力。",
            metadata={"source": "tech_trends.txt", "category": "技术趋势", "topic": "区块链"}
        ),
        Document(
            page_content="《红楼梦》是中国古典四大名著之一，描写了贾宝玉与林黛玉的爱情悲剧。",
            metadata={"source": "classic_literature.txt", "category": "古典文学", "dynasty": "清朝"}
        )
    ]

    dotenv.load_dotenv()

    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL"),
        base_url=os.getenv("CLOSE_BASE_URL"),
        api_key=os.getenv("CLOSE_API_KEY")
    )

    # 获取向量数据库
    db = FAISS.from_documents(documents=documents, embedding=embedding_model)

    # 创建向量搜索器
    retriever = db.as_retriever(search_kwargs={"k": 3}) # 设置返回的文档数量，默认为4

    docs = retriever.invoke("区块链")
    for i, doc in enumerate(docs):
        print(f"{i+1}. {doc.page_content}")

#fun_2()

#分数或阈值查询
def fun_3():
    documents = [
        Document(
            page_content="近日，中国科学家在量子计算领域取得重大突破，实现了100量子比特的纠缠态操控。",
            metadata={"source": "science_news.txt", "category": "科技新闻", "date": "2024-01-15"}
        ),
        Document(
            page_content="《三体》是刘慈欣创作的科幻小说三部曲，包括《三体》、《黑暗森林》和《死神永生》。",
            metadata={"source": "literature_review.txt", "category": "文学作品", "author": "刘慈欣"}
        ),
        Document(
            page_content="人工智能在医疗诊断中的应用越来越广泛，特别是在影像识别方面表现突出。",
            metadata={"source": "ai_research.txt", "category": "学术研究", "field": "医疗AI"}
        ),
        Document(
            page_content="巴黎奥运会筹备工作进入冲刺阶段，预计将有超过200个国家和地区参与。",
            metadata={"source": "sports_news.txt", "category": "体育新闻", "event": "2024巴黎奥运会"}
        ),
        Document(
            page_content="区块链技术不仅应用于加密货币，还在供应链管理、数字身份验证等领域展现潜力。",
            metadata={"source": "tech_trends.txt", "category": "技术趋势", "topic": "区块链"}
        ),
        Document(
            page_content="《红楼梦》是中国古典四大名著之一，描写了贾宝玉与林黛玉的爱情悲剧。",
            metadata={"source": "classic_literature.txt", "category": "古典文学", "dynasty": "清朝"}
        )
    ]

    dotenv.load_dotenv()

    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL"),
        base_url=os.getenv("CLOSE_BASE_URL"),
        api_key=os.getenv("CLOSE_API_KEY")
    )

    # 获取向量数据库
    db = FAISS.from_documents(documents=documents, embedding=embedding_model)

    # 创建向量搜索器
    retriever = db.as_retriever(
        search_type="similarity_score_threshold",
        search_kwargs={"k": 3, "score_threshold": 0.1}  #只有相似度超过这个阈值，才会返回
    )

    docs = retriever.invoke("区块链")
    for i, doc in enumerate(docs):
        print(f"{i+1}. {doc.page_content}")

#fun_3()


#向量检索+LLM模型
def fun_4():
    dotenv.load_dotenv()

    prompt_template = """
        请使用以下提供的文本内容来会啊问题。仅使用提供的文本信息，如果文本中没有相关信息，请回答 ‘抱歉，提供的文本中没有这个信息’
        
        文本内容：{context}
        
        问题：{question}
        
        回答：
    """

    prompt = PromptTemplate.from_template(prompt_template)

    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL"),
        base_url=os.getenv("CLOSE_BASE_URL"),
        api_key=os.getenv("CLOSE_API_KEY")
    )

    # 加载TXT文档
    loader = TextLoader(file_path="F:/data/1.txt", encoding="utf-8")
    documents = loader.load()

    # 创建文本分块器
    test_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=10)
    chunks = test_splitter.split_documents(documents)

    # 获取向量数据库
    db = FAISS.from_documents(documents=documents, embedding=embedding_model)

    retriever = db.as_retriever()

    docs = retriever.invoke("介绍下协和医院的历史")

    chain = prompt | chat_model

    result = chain.invoke(input={"context": docs, "question": "介绍下协和医院历史"})
    print(result)

#fun_4()